#ifndef __UTILS_PE__
#define __UTILS_PE__

namespace xm
{
	struct SECTION_ITEM
	{
		LPCSTR	aSectName;
		DWORD	aUnExpChar;
	};

	static SECTION_ITEM aCompilerSectionNames [] = 
	{
		"CODE\0\0\0\0",		IMAGE_SCN_MEM_WRITE,
		"text\0\0\0\0",		IMAGE_SCN_MEM_WRITE,
		".text\0\0\0",		IMAGE_SCN_MEM_WRITE,
		".itext\0\0",		IMAGE_SCN_MEM_WRITE,
		".data\0\0\0",		IMAGE_SCN_MEM_EXECUTE,
		".rsrc\0\0\0",		IMAGE_SCN_MEM_EXECUTE,
		".reloc\0\0",		IMAGE_SCN_MEM_EXECUTE,
		".idata\0\0",		IMAGE_SCN_MEM_EXECUTE,
		".rdata\0\0",		IMAGE_SCN_MEM_EXECUTE,
		".rodata\0",		IMAGE_SCN_MEM_EXECUTE,
		".ndata\0\0",		IMAGE_SCN_MEM_EXECUTE,
		".sdata\0\0",		IMAGE_SCN_MEM_EXECUTE,
		".bss\0\0\0\0",		IMAGE_SCN_MEM_EXECUTE,
		".crt\0\0\0\0",		IMAGE_SCN_MEM_EXECUTE,
		"INIT\0\0\0\0",		0,
		"PAGE\0\0\0\0",		0,
		".shared\0",		IMAGE_SCN_MEM_EXECUTE,
		".shareda",			IMAGE_SCN_MEM_EXECUTE,
		"RT_CODE\0",		0,
		"RT_DATA\0",		IMAGE_SCN_MEM_EXECUTE,
		"DATA\0\0\0\0",		IMAGE_SCN_MEM_EXECUTE,
		"BSS\0\0\0\0\0",	IMAGE_SCN_MEM_EXECUTE,
		".tls\0\0\0\0",		IMAGE_SCN_MEM_EXECUTE,
		".ndata\0\0",		IMAGE_SCN_MEM_EXECUTE,
		"_DATA\0\0\0",		IMAGE_SCN_MEM_EXECUTE,
		"_BBS\0\0\0\0",		IMAGE_SCN_MEM_EXECUTE,
	};
	enum 
	{
		UNK_SECTION			= 1,
		UNEXP_ATTR			= 2,
	};

	static const char * make_good_section_name( const uint8_t * secName, char * name )
	{
		const uint8_t * p1 = secName;
		size_t cp = 0;
		for( ; cp < 8 && !isspace(p1[cp]); ++ cp );
		if( !cp ) return NULL;
		strncpy( name, (const char*)secName, cp );
		++ cp;
		for( ; cp < 8; ++ cp )
		{
			if( p1[cp] && !isspace(p1[cp]) )
				return NULL;
		}
		return name;
	}

	// MUST 8 Bytes!
	static HRESULT TestCompilerSectionTable( IPeReader * reader, UINT32 & uCase )
	{
		uCase = 0;
		DWORD nSect = 0;
		if( !reader->SectionCount( &nSect ) ) return E_FAIL;
		for( DWORD i = 0; i < nSect; ++ i )
		{
			IPeReader::TImgSecHdr * sections = reader->SectionHeader( i );
			if( !sections ) return E_FAIL;
			int diff = 0;
			for( size_t i = 0; i < sizeof(aCompilerSectionNames)/sizeof(aCompilerSectionNames[0]); ++ i )
			{
				SECTION_ITEM & item = aCompilerSectionNames[i];
				char _name[12] = {};
				const char * snp = make_good_section_name( sections->Name, _name );
				if( !snp ) 
				{
					uCase |= UNK_SECTION;
					continue;
				}
				diff = strnicmp( snp, item.aSectName, 8 );
				if( diff ) continue;
				if( sections->Characteristics & item.aUnExpChar )
				{
					uCase |= UNEXP_ATTR;
					return S_OK;
				}
				break;
			}
			if( diff )
			{
				uCase |= UNK_SECTION;
				return S_OK;
			}
		}
		return S_OK;
	};

	enum { RWX = (IMAGE_SCN_MEM_EXECUTE|IMAGE_SCN_MEM_WRITE) };

	enum 
	{
		ONE_RWX				= 1,
		ONE_MORE_RWX		= 2,
		ENTRY_LAST_1		= 4,
		ENTRY_LAST_N		= 8,
		ENTRY_RWX			= 16,
	};

	static HRESULT TestInfectOrPacked( IPeReader * reader, DWORD & nMask, PDWORD pdwEntrySectinoId = NULL )
	{
		nMask = 0;
		DWORD nRWX = 0, nRVA = 0, nSect = 0;
		//////////////////////////////////////////////////////////////////////////
		if( !reader->SectionCount( &nSect ) ) return E_FAIL;
		if( !reader->EntryPointRva( &nRVA ) ) return E_FAIL;
		//////////////////////////////////////////////////////////////////////////
		IPeReader::TImgDatDir * dde = reader->DataDirectory( WINPE_IMAGE_DIRECTORY_ENTRY_COM_DESCRIPTOR );
		if( dde && dde->VirtualAddress ) return E_FAIL;
		//////////////////////////////////////////////////////////////////////////
		if( nSect > 7 ) return E_FAIL;
		if( !nRVA ) return E_FAIL;
		//////////////////////////////////////////////////////////////////////////
		UINT64 uFileSize = 0;
		if( !reader->FileSize(&uFileSize ) ) return E_FAIL;
		//////////////////////////////////////////////////////////////////////////
		DWORD iEntSec = -1;
		for( DWORD i = 0; i < nSect; ++ i )
		{
			IPeReader::TImgSecHdr * sections = reader->SectionHeader( i );
			if( !sections ) return E_FAIL;
			DWORD Char = sections->Characteristics;
			if( (Char & RWX) == RWX )
			{
				nMask |= ONE_RWX;
				nRWX ++;
			}
			UINT64 nRawEnd = sections->PointerToRawData + sections->SizeOfRawData;
			if( nRawEnd > uFileSize ) return E_FAIL;

			DWORD nSecEnd = sections->VirtualAddress + sections->SizeOfRawData;
			if( nRVA >= sections->VirtualAddress && nRVA < nSecEnd )
			{
				iEntSec = i;
				if( (Char & RWX) == RWX ) nMask |= ENTRY_RWX;
			}
		}
		if( iEntSec != (DWORD)-1 && nSect > 3 )
		{
			if( iEntSec + 1 == nSect )
			{
				nMask |= ENTRY_LAST_1;
			}
			else if( iEntSec + 2 >= nSect )
			{
				nMask |= ENTRY_LAST_N;
			}
		}
		if( pdwEntrySectinoId ) 
			*pdwEntrySectinoId = iEntSec;
		return S_OK;
	}

	static HRESULT TestMSVB( IPeReader * reader )
	{
		DWORD nRVA = 0;
		if( !reader->EntryPointRva( &nRVA ) ) return E_FAIL;

		UINT8 szCode[10] = {};
		DWORD eff = reader->ReadDatRva2( nRVA, szCode, sizeof(szCode) );
		if( eff != sizeof(szCode) ) return E_FAIL;

		if( szCode[0] != 0x68 || szCode[3] != 0x40 || szCode[4] != 0x00 || 
			szCode[5] != 0xE8 || szCode[8] != 0xFF || szCode[9] != 0xFF )
		{
			return E_FAIL;
		}

		UINT32 tagRva = *(le::U32*)(szCode+1);
		if( !reader->ReadDatVa( tagRva, szCode, 4 ) )
			return E_FAIL;

		if( szCode[0] == 'V' && szCode[1] == 'B' && 
			szCode[2] == '5' && szCode[3] == '!' )
		{
			return S_OK;
		}

		return E_FAIL;
	}
	
	static const char * SESSION_XTO_PETOOLS = "this.pe.tools";
	static const char * SESSION_XTO_PETRUNK = "this.pe.trunkfetch";

	static HRESULT get_this_pe_tools( IKDetectSession * sess, IPeReader * wpe, IKPETools ** pptools )
	{
		iptr<IKXoTable> xto = sess->SessionXoTable();
		if( !xto ) return E_FAIL;
		iptr<IKPETools> tool;
		xv::xvalue_t vo;
		if( FAILED( xto->GetXo( SESSION_XTO_PETOOLS, vo ) ) || (!vo.isUnknownPtr()) || !(tool = (IUnknown*)vo) )
		{
			iptr<IKPETools> tools( PROGID_WinPETools );
			if( !tools ) return E_OUTOFMEMORY;
			RFAILED( tools->Init( wpe ) );
			xto->SetUnknown( SESSION_XTO_PETOOLS, tools );
			tool = tools;
		}
		if( !tool ) return E_FAIL;
		*pptools = tool.detach();
		return S_OK;
	}

	static HRESULT get_this_pe_trunk_fetcher( IKDetectSession * sess, IUnknown * ino, IKTrunkFetcher ** pptfo )
	{
		iptr<IKXoTable> xto = sess->SessionXoTable();
		if( !xto ) return E_FAIL;
		iptr<IKTrunkFetcher> trunk;
		xv::xvalue_t vo;
		if( FAILED( xto->GetXo( SESSION_XTO_PETRUNK, vo ) ) || !vo.isUnknownPtr() || !(trunk = (IUnknown*)vo) )
		{
			iptr<IKTrunkFetcher> trunkf( PROGID_WinPEThunkFetcher );
			if( !trunkf ) return E_OUTOFMEMORY;
			RFAILED( trunkf->Open( ino ) );
			xto->SetUnknown( SESSION_XTO_PETRUNK, trunkf );
			trunk = trunkf;
		}
		if( !trunk ) return E_FAIL;
		*pptfo = trunk.detach();
		return S_OK;
	}

};


#endif